import asyncio
import io
import json
import logging
import re
from json import JSONDecodeError
from pathlib import Path
from typing import List, Tuple, Dict

from langchain_community.cache import SQLiteCache
from langchain_community.chat_models import GigaChat
from langchain_community.document_loaders import TextLoader
from langchain_core.documents import Document
from langchain_core.messages import BaseMessage
from langchain_core.prompt_values import ChatPromptValue
from langchain_core.prompts import ChatPromptTemplate
import langchain
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.prompts import load_prompt
from langchain.chains.summarize import load_summarize_chain

from sqlalchemy import inspect


def _sanitalize_text(text: str) -> str:
    text = re.sub(pattern='Украин.', repl='', string=text)
    text = re.sub(pattern='Ukrain.', repl='', string=text)
    return text


def _sanitize_input(request: ChatPromptValue):
    for message in request.messages:
        message.content = _sanitalize_text(message.content)
    return request


class GigaChainWrapper:
    def __init__(
            self, sql_cache_path: Path, giga_token: str, giga_scope: str,
            base_context="Ты ученый в области Computer Science. Ты обязан помогать в написании научных статей и давать ценные советы."
    ):
        self.sql_cache_path = sql_cache_path
        self.giga_token = giga_token
        self.giga_scope = giga_scope
        self.base_context = base_context

        langchain.llm_cache = SQLiteCache(database_path=str(sql_cache_path))

        # Авторизация в сервисе GigaChat
        self.chat = GigaChat(credentials=giga_token, verify_ssl_certs=False, scope=giga_scope)
        self._default_chain = (
                ChatPromptTemplate.from_messages(
                    [
                        ("system", self.base_context),
                        ("human", "{input}")
                    ]
                )
                | _sanitize_input
                | self.chat
        )
        self.logger = logging.getLogger(__name__)

    def _handle_blacklist(self, resp: BaseMessage):
        # TODO possible blacklist processing
        # if ('finish_reason' in resp.response_metadata.keys() and
        #         resp.response_metadata['finish_reason'] in ['blacklist']):
        # inspector = inspect(langchain.llm_cache.engine)
        # schemas = inspector.get_schema_names()
        # for schema in schemas:
        #     print("schema: %s" % schema)
        #     for table_name in inspector.get_table_names(schema=schema):
        #         print(f'table_name = {table_name}')
        #         for column in inspector.get_columns(table_name, schema=schema):
        #             print("Column: %s" % column)
        #         with langchain.llm_cache.engine.connect() as connection:
        #             result = connection.execute(text(f"select * from {table_name};"))
        #             for row in result:
        #                 print(result)
        pass

    async def get_search_queries(self, text: str, n_queries: int) -> List[str]:
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.base_context),
                ("human", f"Придумай {n_queries} различных поисковых запросов на английском языке,"
                          "которые помогут найти справочную литературу по следующей теме: \n\n{input}\n\nРезультаты представь в формате json списка")

            ]
        )
        chain = (
                prompt
                | _sanitize_input
                | self.chat
        )
        resp = await chain.ainvoke({'input': text})
        try:
            l = json.loads(resp.content)
            assert isinstance(l, list)
            # assert len(l) == len(set(l)), f"{len(l)} != {len(set(l))}"
            return list(set(l))
        except JSONDecodeError as e:
            return None

    async def get_summary_book(self, text: str) -> BaseMessage:
        documents = [Document(text)]
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=7000,
            chunk_overlap=0,
            length_function=len,
            is_separator_regex=False,
        )
        documents = text_splitter.split_documents(documents)

        book_map_prompt = load_prompt("lc://prompts/summarize/map_reduce/summarize_book_map.yaml")
        book_combine_prompt = load_prompt("lc://prompts/summarize/map_reduce/summarize_book_combine.yaml")

        chain = load_summarize_chain(
            self.chat, chain_type="map_reduce",
            map_prompt=book_map_prompt,
            combine_prompt=book_combine_prompt,
            verbose=False
        )

        resp = await chain.ainvoke({"input_documents": documents})

        return resp

    # use this for summary of article annotation
    async def get_summary_simple(self, text: str) -> List[str]:
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.base_context),
                ("human", "Сделай краткий пересказ на русском языке приведенного далее текста в 2-3 предложениях:"
                          "\n\n{input}\n\n. ")
            ]
        )
        chain = (
                prompt
                | _sanitize_input
                | self.chat
        )
        resp = await chain.ainvoke({'input': text})

        self.logger.info(resp)

        return resp.content

    async def get_annotation(self, description: str, summary_list: List[Tuple[str, str]]) -> str:
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.base_context),
                ("human", "Твой коллега хочет обсудить с тобой следующую идею:"
                          "\n\n{input}\n\n "),
                ("human", "Он нашел научные статьи по предложенной идее.")
            ]
            +
            [("human", f'В статье {name} описано следующее: \n\n{desc}\n\n') for name, desc in summary_list]
            +
            [("human",
              "Помоги ему составить описание для собственной научной статьи, дополни его своими идеями. Напиши до 10 предложений.")]
        )
        chain = (
                prompt
                | _sanitize_input
                | self.chat
        )
        resp = await chain.ainvoke({'input': description})
        return resp.content

    async def get_themes(self, description: str, summary_list: List[Tuple[str, str]], n_themes: int) -> List[str]:
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.base_context),
                ("human", "Твой коллега хочет обсудить с тобой следующую идею:"
                          "\n\n{input}\n\n "),
                ("human", "Он нашел научные статьи по предложенной идее.")
            ]
            +
            [("human", f'В статье {name} описано следующее: \n\n{desc}\n\n') for name, desc in summary_list]
            +
            [("human",
              f"Выдели {n_themes} штук самых главных тем из предложенного материала. Они должны быть полезны для идеи коллеги. Отдели их переводом строки.")]
        )
        chain = (
                prompt
                | _sanitize_input
                | self.chat
        )
        resp = await chain.ainvoke({'input': description})
        return resp.content.split('\n')

    # in summary list must be fields 'title', 'summary', 'bibtex'
    async def get_rel_works(self, description: str, summary_list: List[Dict[str, str]]) -> str:
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.base_context),
                ("human", "Твой коллега хочет обсудить с тобой следующую идею научной статьи:"
                          "\n\n{input}\n\n "),
                ("human", "Он подобрал список литературы, по которому ты должен составить литературный обзор.\n\n")
            ]
            +
            [("human", f"В статье {d['title']} описано следующее:\n"
                       f"{d['summary']}\n"
                       f"Для нее есть bibtex библиография: \n{d['bibtex'].replace('{', '{{').replace('}', '}}')}\n\n")
             for d in summary_list]
            +
            [("human", f"Составь литературный обзор. Сделай ссылки на литературу.")]
        )
        chain = (
                prompt
                | _sanitize_input
                | self.chat
        )
        resp = await chain.ainvoke({'input': description})
        resp_str = resp.content
        return resp_str

    async def get_summary_for_topic(self, topic: str, summary_list: List[Tuple[str, str]]) -> str:
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.base_context),
                ("human", "Твой коллега хочет обсудить с тобой следующую тему для научной статьи:"
                          "\n\n{input}\n\n "),
                ("human", "В качестве справочного материала тебе предлагаются следующие научные статьи.")
            ]
            +
            [("human", f'В статье {name} описано следующее: \n\n{desc}\n\n') for name, desc in summary_list]
            +
            [("human", f"Напиши свои мысли по этой теме, опираясь на справочный материал.")]
        )
        chain = (
                prompt
                | _sanitize_input
                | self.chat
        )
        resp = await chain.ainvoke({'input': topic})
        return resp.content

    # summary_list: list of dicts with fields title and summary
    # output is generator of dicts with fields title and summary
    async def get_generator_for_topics(self, idea: str, summary_list: List[Dict[str, str]], n_topics: int):
        list_of_content = []
        for d in summary_list:
            list_of_content.append((d['title'], d['summary']))
        list_of_themes = await self.get_themes(idea, list_of_content, n_topics)
        for theme in list_of_themes:
            summary = await self.get_summary_for_topic(theme, list_of_content)
            yield {'title': theme, 'summary': summary}, len(list_of_themes)


async def main_test():
    gchain_wrapper = GigaChainWrapper(
        sql_cache_path=Path('sqlite_cache.db'),
        giga_token='Yjg4MTQzMmUtNDAwMS00NDk0LThjOGUtNmU5ZWQ2YzQ4NDQ2OmQ4MWMxZGZiLTFmNGYtNDk5NS05OGQzLTBiMzYyYWJmNjk3OA==',
        giga_scope='GIGACHAT_API_CORP'
    )
    quieries = await gchain_wrapper.get_search_queries('Деревья решений', n_queries=3)
    print(quieries)
    summary = await gchain_wrapper.get_summary_simple(
        """
        Дерево принятия решений (также называют деревом классификации или регрессионным деревом) — средство поддержки принятия решений, использующееся в машинном обучении, анализе данных и статистике. Структура дерева представляет собой «листья» и «ветки». На рёбрах («ветках») дерева решения записаны признаки, от которых зависит целевая функция, в «листьях» записаны значения целевой функции, а в остальных узлах — признаки, по которым различаются случаи. Чтобы классифицировать новый случай, надо спуститься по дереву до листа и выдать соответствующее значение.

Подобные деревья решений широко используются в интеллектуальном анализе данных. Цель состоит в том, чтобы создать модель, которая предсказывает значение целевой переменной на основе нескольких переменных на входе.


Каждый лист представляет собой значение целевой переменной, изменённой в ходе движения от корня по рёбрам дерева до листа. Каждый внутренний узел сопоставляется с одной из входных переменных.

Дерево может быть также «изучено» разделением исходных наборов переменных на подмножества, основанные на проверке значений признаков. Это действие повторяется на каждом из полученных подмножеств. Рекурсия завершается тогда, когда подмножество в узле имеет те же значения целевой переменной, таким образом, оно не добавляет ценности для предсказаний. Процесс, идущий «сверху вниз», индукция деревьев решений (TDIDT)[1], является примером поглощающего «жадного» алгоритма, и на сегодняшний день является наиболее распространённой стратегией деревьев решений для данных, но это не единственная возможная стратегия.

В интеллектуальном анализе данных, деревья решений могут быть использованы в качестве математических и вычислительных методов, чтобы помочь описать, классифицировать и обобщить набор данных, которые могут быть записаны следующим образом:
        """
    )
    print(summary)


if __name__ == '__main__':
    asyncio.run(main_test())
